package no.awi.wpd.probability;

import java.math.BigDecimal;
import java.math.MathContext;
import java.math.RoundingMode;
import java.util.Arrays;

/**
 * A discrete numerical probability distribution, mapping number values to
 * probabilities. Some examples:
 * <ul>
 * <li> A 6 sided dice: <br>
 * {1: 0.166, 2: 0.166, 3: 0.166, 4: 0.166, 5: 0.166, 6: 0.166} </li>
 * <li> Number of heads when tossing 4 coins: <br>
 * {0: 0.0625, 1: 0.25, 2: 0.375, 3: 0.25, 4: 0.0625} </li>
 * </ul>
 */
public class ProbabilityDistribution {

    private static final ProbabilityDistribution ZERO = createUniformProbabilityDistribution(1, 0);

    private static final ProbabilityDistribution ONE = createUniformProbabilityDistribution(1, 1);

    private static final double EPSILON = .0000001;

    /**
     * E.g. a series of coin tosses.
     * 
     * See <a href="http://en.wikipedia.org/wiki/Binomial_distribution">Wikipedia Binomial Distribution</a> for an explanation.
     * 
     * @param p -
     *            the probability of the single event
     * @param n -
     *            the size of the probability distribution
     * @return
     */
    public static ProbabilityDistribution createBinomialProbabilityDistribution(double p, int n) {
    	return createBinomialProbabilityDistribution(p, n, 32);
    }

    /**
     * E.g. a series of coin tosses.
     * 
     * See <a href="http://en.wikipedia.org/wiki/Binomial_distribution">Wikipedia Binomial Distribution</a> for an explanation.
     * 
     * @param p -
     *            the probability of the single event
     * @param n -
     *            the size of the probability distribution
     * @param precision - the precision level used in calculations
     * @return
     */
    public static ProbabilityDistribution createBinomialProbabilityDistribution(double p, int n, int precision) {
        if (n < 0 || p < 0 || p > 1) {
            throw new IllegalArgumentException();
        }
        if (n == 0) {
            return ZERO();
        }
        MathContext mc = new MathContext(precision);
        ProbabilityDistribution binomialProbabilityDistribution = new ProbabilityDistribution(new double[n + 1]);
        BigDecimal binomialCoefficient = BigDecimal.valueOf(1);
        BigDecimal denominator = BigDecimal.valueOf(1);
        BigDecimal numerator = BigDecimal.valueOf(1);
        for (int i = 0; i < n + 1; i++) {
			double result = nextProbability(p, n, binomialCoefficient, i, mc);
			binomialProbabilityDistribution.setProbabilityByIndex(i, result);
			
			
			numerator = nextNumerator(n, numerator, i, mc);
			denominator = nextDenominator(denominator, i, mc);
			binomialCoefficient = nextBinomialCoefficient(denominator, numerator, mc);
        }
        return binomialProbabilityDistribution;
    }

	private static double nextProbability(double p, int n,
			BigDecimal binomialCoefficient, int i, MathContext mc) {
		BigDecimal positiveProbability = BigDecimal.valueOf(p).pow(i, mc);
		BigDecimal negationProbability = BigDecimal.valueOf(1 - p).pow(n - i, mc);
		double result = positiveProbability.multiply(negationProbability, mc).multiply(binomialCoefficient, mc).doubleValue();
		return result;
	}

	private static BigDecimal nextBinomialCoefficient(BigDecimal denominator,
			BigDecimal numerator, MathContext mc) {
		BigDecimal binomialCoefficient;
		binomialCoefficient = numerator.divide(denominator, mc);
		return binomialCoefficient;
	}

	private static BigDecimal nextDenominator(BigDecimal denominator, int i, MathContext mc) {
		denominator = denominator.multiply(BigDecimal.valueOf(i + 1), mc);
		return denominator;
	}

	private static BigDecimal nextNumerator(int n, BigDecimal numerator, int i, MathContext mc) {
		numerator = numerator.multiply(BigDecimal.valueOf(n - i), mc);
		return numerator;
	}

    /**
     * Creates {@link ProbabilityDistribution} that has the same probability for every
     * value (and sum=1.0 of course).
     * 
     * @param size
     * @param minValue
     * @return
     */
    public static ProbabilityDistribution createUniformProbabilityDistribution(int size, int minValue) {
        ProbabilityDistribution evenProbabilityDistribution = new ProbabilityDistribution(new double[size]);
        evenProbabilityDistribution.minValue = minValue;
        double probability = 1.0 / size;
        for (int i = 0; i < size; i++) {
            evenProbabilityDistribution.setProbabilityByIndex(i, probability);
        }
        return evenProbabilityDistribution;
    }

    private double[] distribution;

    private int minValue;

    public ProbabilityDistribution() {
        /* Empty */
    }

    public ProbabilityDistribution(int possibilities) {
        this(possibilities, 0);
    }

    public ProbabilityDistribution(int possibilities, int minValue) {
        this(new double[possibilities], minValue);
    }

    public ProbabilityDistribution(double[] distribution) {
        this(distribution, 0);
    }

    public ProbabilityDistribution(double[] distribution, int minValue) {
        this.distribution = distribution;
        this.minValue = minValue;
    }

    public static ProbabilityDistribution ONE() {
        return ONE.clone();
    }

    /**
     * Given two ProbabilityDistributions:
     * <ul>
     * <li>{0: 0.5, 1: 0.5}</li>
     * <li>{0: 0.25, 1: 0.5, 2: 0.25}</li>
     * </ul>
     * The result of adding them absolutely is
     * <p> 
     * {0: 0.75, 1: 1.0, 2: 0.25}
     * <p>
     * 
     * This method is symmetric (a.addAbsolutely(b) === b.addAbsolutely(a)).
     * 
     * @param other
     * @return a new {@link ProbabilityDistribution} instance
     */
    public ProbabilityDistribution addAbsolutely(ProbabilityDistribution other) {
        ProbabilityDistribution result = new ProbabilityDistribution(new double[Math.max(getMaxValue(), other.getMaxValue()) + 1]);
        result.minValue = Math.min(minValue, other.minValue);
        for (int i = 0; i < result.getSize(); i++) {
	    int value = result.getValueByIndex(i);
	    double newProbability = 0;
	    if (this.hasValue(value))
		newProbability += this.getProbabilityByValue(value);
	    if (other.hasValue(value)) 
		newProbability += other.getProbabilityByValue(value);
	    result.setProbabilityByValue(value, newProbability);
        }
        return result;
    }

    /**
     * Adds two ProbabilityDistributions together by value. That is, given two
     * ProbabilityDistributions representing a 6 sided dice each, the result of
     * joining them is the {@link ProbabilityDistribution} representing 2D6.
     * 
     * @param other
     * @return a new {@link ProbabilityDistribution} instance
     */
    public ProbabilityDistribution add(ProbabilityDistribution other) {
        int resultMinValue = minValue + other.minValue;
        ProbabilityDistribution result = new ProbabilityDistribution(new double[getMaxValue() + other.getMaxValue()
                - resultMinValue + 1], resultMinValue);
        for (int i = 0; i < getSize(); i++) {
            for (int j = 0; j < other.getSize(); j++) {
                double probabilityDelta = getProbabilityByIndex(i) * other.getProbabilityByIndex(j);
                int newValue = getValueByIndex(i) + other.getValueByIndex(j);
                result.addProbabilityByIndex(result.getIndexByValue(newValue), probabilityDelta);
            }
        }
        return result;
    }
    
    /**
     * All probabilities are multiplied by the given value.
     * <p>
     * Note that this method will change the total sum of probabilities. 
     * 
     * @param constant
     */
    public void multiplyByConstant(double constant) {
        for (int i = 0; i < getSize(); i++) {
            getDistribution()[i] *= constant;
        }
    }

    private double[] getDistribution() {
        return distribution;
    }

    public int getMaxValue() {
        return minValue + getSize() - 1;
    }

    public int getMinValue() {
        return minValue;
    }

    public void setMinValue(int minValue) {
        this.minValue = minValue;
    }

    public double getProbabilityByValue(int value) {
        return getDistribution()[getIndexByValue(value)];
    }
    
    public double getProbabilityByIndex(int index) {
        return getDistribution()[index];
    }

    public void addProbabilityByIndex(int index, double probabilityDelta) {
        setProbabilityByIndex(index, getProbabilityByIndex(index) + probabilityDelta);
    }
    
    public void addProbabilityByValue(int value, double probabilityDelta) {
        addProbabilityByIndex(getIndexByValue(value), probabilityDelta);
    }

    public double getSum() {
    	return getSumOverTreshhold(minValue);
    }

    /**
     * Given a {@link ProbabilityDistribution}:
     * <p>
     * {0: 0.25, 1: 0.25, 2: 0.25, 3: 0.25}
     * <p>
     * {@link #getSumOverTreshhold(2)} will return 0.5
     * 
     * @param treshhold - the lower limit value (inclusive)  
     */
    public double getSumOverTreshhold(int treshhold) {
        int index = getIndexByValue(treshhold);
        if (index < 0 || index >= getSize()) {
            throw new IllegalArgumentException();
        }
        double sum = 0;
        for (int i = index; i < getSize(); i++) {
            sum += getProbabilityByIndex(i);
        }
        return sum;
    }

    public ProbabilityDistribution withMaxValue(int maxValue) {
        if (maxValue >= getMaxValue()) {
            return this.clone();
        }
        ProbabilityDistribution result = new ProbabilityDistribution(maxValue + Math.abs(minValue) + 1, minValue);
        for (int index = 0; index < result.distribution.length; index++) {
           if (getValueByIndex(index) == maxValue) {
               result.setProbabilityByIndex(index, getSumOverTreshhold(maxValue));
               return result;
           }
           result.distribution[index] = getProbabilityByIndex(index);
        }
        return result;
    }

    /**
     * Given a {@link ProbabilityDistribution}:
     * <p>
     * {0: 0.25, 1: 0.25, 2: 0.25, 3: 0.25}
     * <p>
     * {@link #getSumUnderTreshhold(2))} will return 0.75
     * 
     * @param treshhold - the upper limit value (inclusive)  
     */
    public double getSumUnderTreshhold(int treshhold) {
        int index = getIndexByValue(treshhold);
        if (index < 0 || index >= getSize()) {
            throw new IllegalArgumentException();
        }
        double sum = 0;
        for (int i = index; i > -1; i--) {
            sum += getProbabilityByIndex(i);
        }
        return sum;
    }
    
    public int getValueByIndex(int index) {
        return minValue + index;
    }

    public int getIndexByValue(int value) {
        return value - minValue;
    }

    public boolean hasValue(int value) {
        return getIndexByValue(value) > -1 && getIndexByValue(value) < distribution.length;
    }

    public void setProbabilityByIndex(int index, double value) {
        distribution[index] = value;
    }

    public void setProbabilityByValue(int value, double probability) {
        distribution[getIndexByValue(value)] = probability;
    }

    public int getSize() {
        return distribution.length;
    }
    
    /**
     * Validates that the total sum of probabilities is 1.0 +- 0.0000005
     * 
     */
    public boolean validateSum() {
    	double sum = getSum();
		return sum < 1 + EPSILON && sum > 1 - EPSILON;
    }
    
    @Override
    public String toString() {
        String result = "{";
        for (int i = 0; i < getSize(); i++) {
            result += getValueByIndex(i) + ": " + String.format("%.1f%%", getDistribution()[i] * 100);
            if (i != getSize() - 1) {
                result += ", ";
            }
        }
        return result + "} = " + getSum();
    }

    public String confidenceInterval() {
	    if (!validateSum())
            throw new IllegalStateException();
	    double maxProbability = 0;
	    int maxValueIndex = 0;
	    for (int i = 0; i < getSize(); i++) {
	        if (distribution[i] > maxProbability) {
                maxProbability = distribution[i];
                maxValueIndex = i;
	        }
	    }
        double treshhold = 0.75;
        double sumProbabilities = maxProbability;
        int rightSideIndex = maxValueIndex + 1;
        int leftSideIndex = maxValueIndex - 1;
        while (sumProbabilities < treshhold) {
            if (leftSideIndex > 0) {
            sumProbabilities += getProbabilityByIndex(leftSideIndex);
            }
            if (rightSideIndex < distribution.length) {
            sumProbabilities += getProbabilityByIndex(rightSideIndex);
            }
            if (sumProbabilities > treshhold) break;
            rightSideIndex = Math.min(distribution.length - 1, rightSideIndex + 1);
            leftSideIndex = Math.max(0, leftSideIndex - 1);
        }
        leftSideIndex = Math.max(0, leftSideIndex);
        rightSideIndex = Math.min(distribution.length - 1, rightSideIndex);
        return "[" + getValueByIndex(maxValueIndex) + " +" + (rightSideIndex - maxValueIndex) + "/-" +
                (maxValueIndex - leftSideIndex) + " " + String.format("%.1f%%", sumProbabilities * 100) + "]";
    }

    public String toJson() {
    	String result = "[";
    	for (int i = 0; i < getSize(); i++) {
    		result += "{ value: " + getValueByIndex(i) + ", probability: " + getDistribution()[i] + "}";
    		if (i != getSize() - 1) {
    			result += ", ";
    		}
    	}
    	return result + "]";
    }
    
    @Override
    public int hashCode() {
        final int prime = 31;
        int result = 1;
        result = prime * result + Arrays.hashCode(distribution);
        result = prime * result + minValue;
        return result;
    }

    @Override
    public boolean equals(Object obj) {
        if (this == obj)
            return true;
        if (obj == null)
            return false;
        if (getClass() != obj.getClass())
            return false;
        final ProbabilityDistribution other = (ProbabilityDistribution) obj;
        if (minValue != other.minValue) {
            return false;
        }
        if (distribution.length != other.distribution.length) {
            return false;
        }
        for (int i = 0; i < distribution.length; i++) {
            if (Math.abs(distribution[i] - other.distribution[i]) > EPSILON) {
                return false;
            }
        }
        return true;
    }

    public ProbabilityDistribution clone() {
        ProbabilityDistribution clone = new ProbabilityDistribution(new double[getSize()]);
        for (int i = 0; i < getSize(); i++) {
            clone.setProbabilityByIndex(i, getProbabilityByIndex(i));
        }
        clone.minValue = minValue;
        return clone;
    }

    /**
     * @return - a {@link ProbabilityDistribution} with 100% chance of 0
     */
    public static ProbabilityDistribution ZERO() {
        return ZERO.clone();
    }

    /**
	 * 
	 * 
	 * @param subtrahend
	 * @return - a new {@link ProbabilityDistribution} object that is the difference between this and subtrahend
	 */
	public ProbabilityDistribution subtract(ProbabilityDistribution subtrahend) {
		ProbabilityDistribution difference = new ProbabilityDistribution(getSize() + subtrahend.getSize() - 1);
		difference.setMinValue(getMinValue() - subtrahend.getMaxValue());
		
		for (int i = 0; i < getDistribution().length; i++) {
			for (int j = 0; j < subtrahend.getDistribution().length; j++) {
				double delta = getProbabilityByIndex(i) * subtrahend.getProbabilityByIndex(j);
				difference.addProbabilityByIndex(i + subtrahend.getSize() - 1 - j, delta);
			}
		}
		return difference;
	}

	/**
	 * Remove any trailing values with probability = 0 at both ends of the
	 * distribution
	 */
	public void trim() {
		int length = distribution.length;
		int minValue = this.minValue;
		int lowerLimit = 0;
		for (; lowerLimit < distribution.length; lowerLimit++) {
			if (distribution[lowerLimit] == 0) {
				minValue++;
				length--;
			} else {
				break;
			}
		}
		for (int i = distribution.length - 1; i >= 0; i--) {
			if (distribution[i] == 0 && length > 0) {
				length--;
			} else {
				break;
			}
		}
		double[] distribution = new double[length];
		for (int i = 0; i < distribution.length; i++) {
			distribution[i] = this.distribution[lowerLimit];
			lowerLimit++;
		}
		this.distribution = distribution;
		this.minValue = minValue;
	}

	/**
	 * 
	 * 
	 * @return
	 */
	public ProbabilityDistribution negateAndReverse() {
		ProbabilityDistribution result = new ProbabilityDistribution(getSize(), (- this.getMaxValue()));
		for (int i = 0; i < result.getSize(); i++) {
			result.setProbabilityByIndex(i, getProbabilityByIndex(getSize() - i - 1));
		}
		return result;
	}


    public ProbabilityDistribution multiplyValuesByConstant(int multiplier) {
        if (multiplier == 0) {
            return ZERO();
        }
        if (multiplier < 0) {
            throw new IllegalArgumentException("Illegal negative multiplier: " + multiplier);
        }
        ProbabilityDistribution result = new ProbabilityDistribution(this.distribution.length * multiplier,
                minValue * multiplier);
        for (int i = 0 ; i < this.distribution.length; i++) {
            result.setProbabilityByValue(this.getValueByIndex(i) * multiplier, this.getProbabilityByIndex(i));
        }
        return result;
    }

    public ProbabilityDistribution conditionalOrValue(double conditional, int value) {
        if (!hasValue(value)) throw new IllegalArgumentException();
        ProbabilityDistribution result = clone();
        result.setProbabilityByValue(value, 1 - conditional * (1 - result.getProbabilityByValue(value)));
        for (int i = 1; i < result.distribution.length; i++) {
            result.setProbabilityByValue(i, conditional * result.getProbabilityByValue(i));
        }
        return result;
    }

    public ProbabilityDistribution multiply(int n) {
        if (n < 0) throw new IllegalArgumentException("Cannot multiply ProbabilityDistribution by negative number");
        if (n == 0) {
            return ZERO();
        }
        ProbabilityDistribution result = this.clone();
        for (int i = 1; i < n; i++) {
            result = this.add(result);
        }
        return result;
    }
}
